import torch
import torch.nn as nn
import torch.nn.functional as F

class SelfAttention(nn.Module):
    def __init__(self, hidden_size):
        super(SelfAttention, self).__init__()
        self.attention = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.Tanh(),
            nn.Linear(hidden_size, 1)
        )
    
    def forward(self, x):
        attention_weights = F.softmax(self.attention(x), dim=1)
        return torch.sum(attention_weights * x, dim=1)

class SentimentClassifier(nn.Module):
    def __init__(self, n_classes, pretrained_model):
        super(SentimentClassifier, self).__init__()
        self.bert = pretrained_model
        self.hidden_size = self.bert.config.hidden_size
        
        # 注意力层
        self.attention = SelfAttention(self.hidden_size)
        
        # 主干网络
        self.drop1 = nn.Dropout(p=0.4)
        self.fc1 = nn.Linear(self.hidden_size, 512)
        self.norm1 = nn.LayerNorm(512)
        
        self.drop2 = nn.Dropout(p=0.3)
        self.fc2 = nn.Linear(512, 256)
        self.norm2 = nn.LayerNorm(256)
        
        self.drop3 = nn.Dropout(p=0.2)
        self.fc3 = nn.Linear(256, 128)
        self.norm3 = nn.LayerNorm(128)
        
        # 残差连接的投影层
        self.projection1 = nn.Linear(512, 256)
        self.projection2 = nn.Linear(256, 128)
        
        # 输出层
        self.fc4 = nn.Linear(128, n_classes)
        
        # 激活函数
        self.gelu = nn.GELU()
        
    def forward(self, input_ids, attention_mask):
        # BERT编码
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask
        )
        
        # 获取所有token的输出
        sequence_output = outputs[0]
        
        # 应用自注意力机制
        attended = self.attention(sequence_output)
        
        # 主干网络处理
        x = self.drop1(attended)
        x = self.fc1(x)
        x = self.norm1(x)
        x = self.gelu(x)
        
        # 第一层残差连接
        residual = self.projection1(x)
        x = self.drop2(x)
        x = self.fc2(x)
        x = self.norm2(x)
        x = self.gelu(x + residual)
        
        # 第二层残差连接
        residual = self.projection2(x)
        x = self.drop3(x)
        x = self.fc3(x)
        x = self.norm3(x)
        x = self.gelu(x + residual)
        
        return self.fc4(x)
